##########################################################################################

library(ComplexHeatmap)
library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(circlize)

##########################################################################################

option_list <- list(
    make_option(c("--images_path"), type = "character") ,
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--class_order_file"), type = "character") ,
    make_option(c("--class_order_sub_file"), type = "character") ,
    make_option(c("--ccf_file"), type = "character"),
    make_option(c("--tp53_pre_file"), type = "character"),
    make_option(c("--tp53_cancer_file"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"
    ccf_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.addShare.tsv"
	images_path <- paste(work_dir,"/","images",sep="")
	info_file <- paste(work_dir,"/config/tumor_normal.class.list",sep="")
	class_order_file <- paste(work_dir,"/config/Class_order.list",sep="")
	class_order_sub_file <- paste(work_dir,"/config/Class_order_sub.list",sep="")

	tp53_pre_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/lollipop/TP53_NMU/TP53.PreCancerous.UniqueNormal.tsv"
 	tp53_cancer_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/lollipop/TP53_NMU/TP53.Cancerous.UniqueNormal.tsv"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

info_file <- opt$info_file
images_path <- opt$images_path
class_order_file <- opt$class_order_file
class_order_sub_file <- opt$class_order_sub_file
tp53_pre_file <- opt$tp53_pre_file
tp53_cancer_file <- opt$tp53_cancer_file
ccf_file <- opt$ccf_file

###########################################################################################
#smg <- data.frame(fread(smg_list , header = F))
dat_ccf <- data.frame(fread(ccf_file))

class_order <- data.frame(fread(class_order_file , header = T))
class_order_sub <- data.frame(fread(class_order_sub_file , header = T))

info <- data.frame(fread(info_file))

## 用于样本的排序
dat_tp53_pre <- fread(tp53_pre_file)
dat_tp53_can <- fread(tp53_cancer_file)

###########################################################################################
## 对样本进行排序
info$Class <- factor(info$Class,levels= unique(class_order$Class), ordered=TRUE)
info$Class_sub <- factor(info$Class_sub,levels= unique(class_order_sub$Class), ordered=TRUE)

info$ID_order <- paste0(info$ID,"_", as.numeric(info$Class_sub),"_",info$Class)
info <- info[order(info$ID_order),]

###########################################################################################
## show_gene
evolution_gene <- c("DNAH3" , "MICAL2" , "NRG1")
show_gene <- c("TP53" , "DNAH3" , "MICAL2" , "NRG1")
Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

## "TP53" , "APC" , "PIK3CA" , "CDH1"的Trunk突变样本
dat_ccf <- subset( dat_ccf , Hugo_Symbol %in% show_gene & Variant_Classification %in% Variant_Types  )
use_sample <- subset( dat_ccf , Hugo_Symbol %in% evolution_gene )$ID

## 存在3个基因突变样本以及IM + IGC + DGC均有的样本
info <- subset( info , ID %in% use_sample | Type == "IM + IGC + DGC" )
use_sample <- unique(info$ID)
dat_ccf <- subset( dat_ccf , ID %in% use_sample  )
show_gene <- unique(dat_ccf$Hugo_Symbol)

###########################################################################################
## 先按照TP53排序
if(1!=1){
	normal_order1 <- unique(info[info$Normal %in% dat_tp53_pre$Normal,"ID"])
	normal_order2 <- unique(info[info$Normal %in% dat_tp53_can$Normal,"ID"])
	normal_order2 <-  normal_order2[!(normal_order2 %in% normal_order1)]
	normal_order_tp53 <- c( normal_order1 , normal_order2 )
	## order的顺序
	tp53_igc <- c("JZGC00546" , "JZGC00638" , "JZGC00744" , "JZGC00731" ,
		"JZGC00664" , "JZGC00905" , "JZGC01075")
	tp53_dgc <- c("JZGC00751" , "JZGC00553" , "JZGC00549")
	apc_igc <- c("JZGC00570")
	apc_dgc <- c("JZGC00941")
	pik3ca_igc <- c("JZGC01033")
	pik3ca_dgc <- c("JZGC00464")
	cdh1_dgc <- c("JZGC00618")
	normal_order <- c(
		tp53_igc , apc_igc , pik3ca_igc , 
		tp53_dgc , apc_dgc , pik3ca_dgc , cdh1_dgc
		)

	info$ID <- factor(info$ID,levels= normal_order, ordered=TRUE)
	info <- info[order(info$ID , info$ID_order),]
}

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")


col_class <- col[1:3]

CreateMutMatrix <- function(dat = dat , Variant_Type = Variant_Type ){

	pre <- c("IM")
	can <- c("IGC","DGC")

	mut <- dat
	mut$Variant_Classification <- as.character(mut$Variant_Classification)
	mut <- mut[which(mut$Variant_Classification %in% Variant_Type),]
	mut$use_share <- ifelse( mut$Share == "TRUE" , "Share" , "Private" )
	mut$clonal_status <- ifelse( mut$CCF_adj >= 0.6 , "Clonal" , "Subclonal" )
	mut$cnv_status <- ifelse( mut$total_cn==2 & mut$minor_cn==0 , "LOH" , "" )
	mut$cnv_status <- ifelse( mut$total_cn > 2 , "AMP" , mut$cnv_status)
	mut$cnv_status <- ifelse( mut$total_cn == 1 & mut$minor_cn==0 , "LOSS" , mut$cnv_status )

	## INDEL 和 INS 合并
	mut[grep("In_Frame",mut$Variant_Classification),'Variant_Classification'] = "In_Frame"
	mut[grep("Frame_Shift",mut$Variant_Classification),'Variant_Classification'] = "Frame_Shift"

	## Gene只考虑有突变的SMG和CGC
	Gene <- show_gene

	## 构建矩阵
	Sample <- info[,"Tumor"]
	maf_matrix <- matrix("" , ncol = length(Sample) , nrow = length(Gene) , dimnames = list(Gene,Sample))

	for(gene in rownames(maf_matrix)){
		print(gene)

		## 克隆类型
		for(tumor in colnames(maf_matrix)){
			index <- which(mut$Hugo_Symbol==gene & mut$Tumor==tumor)
			if(length(index)==0){
				var=""
			}else if(length(index)==1){
				var <- mut[index,'clonal_status']
				#var <- paste(mut[index,'clonal_status'] , mut[index,'cnv_status'] , sep = ";")

			}else if(length(index)>1){
				if(length( which(mut[index,'clonal_status'] == "Clonal") != 0 )){
					var <- unique(mut[index,'clonal_status'])
					#var <- unique( paste( mut[index,'clonal_status'] , mut[index,'cnv_status'] , sep = ";") )
				}else{
					var <- "Subclonal"
				}
			}
			print(var)
			maf_matrix[ which(rownames(maf_matrix)==gene) , which(colnames(maf_matrix)==tumor) ] <-  var
		}
	}

	return(maf_matrix)
}

MutMatrixOrder <- function( mut = mut , info = info){

	######################################
	## 基因排序
	## 按照患肿瘤的病人，然后按照癌前病变有的病人
	IGC_sample <- info[info$Class=="IGC","Tumor"]
	DGC_sample <- info[info$Class=="DGC","Tumor"]
	IM_sample <- info[info$Class=="IM","Tumor"]

	mut_tumor <- mut[,colnames(mut) %in% c(IGC_sample,DGC_sample)]
	mut_pre <- mut[,colnames(mut) %in% c(IM_sample)]

	NumMut_Tumor <- apply(mut_tumor , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	NumMut_Pre <- apply(mut_pre , 1 ,function(x){
		length(unique(info[info$Tumor %in% names(which(x!="")) ,"Normal"]))}
	)

	MutNumMatrix <- data.frame(cbind(NumMut_Pre,NumMut_Tumor))

	## 基因排序
	mut <- mut[order(MutNumMatrix$NumMut_Tumor , MutNumMatrix$NumMut_Pre , decreasing = T ),]

	return(mut)
}

plotMutWaterFull <- function( MutMatrix = MutMatrix , col_class = col_class , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name ){

	################################################################################################
	## 注释的名字
	annotation_name <- c("Class")

	## 展示出现频率 >= 1个样本的突变
	mut <- MutMatrix[apply(MutMatrix,1,function(x){length(which(x!=""))>=1}),]

	## 突变矩阵排序
	mut <- MutMatrixOrder( mut = mut , info = info )
	
	## 顺序,很重要！
	class_order <- as.character(info$Class)
	sample_order <- colnames(mut)
	type_order <- as.character(info$Type)
	type_order <- factor(type_order , levels = c("IM + IGC" , "IM + DGC" , "IM + IGC + DGC") , order = T)

	################################################################################################
	## 突变的颜色
	col = c(
		"orange" , rgb(red=90,green=147,blue=189,alpha=255,max=255) ,
		"white")

	names(col) = c(
	  'Clonal',
	  'Subclonal',
	  'Private'
	)

	col_cnv <-c( "#00A087FF" , "#E64B35FF" , "#3C5488FF" )
	names(col_cnv) <- c("LOH" ,
	  "AMP" ,
	  "LOSS")

	## 设置不同背景的颜色
	if(1!=1){
		alter_fun = list(
		    background = function(x, y, w, h) {
		        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
		            gp = gpar(fill = "#F2F2F2", col = NA))
		    },
		    Clonal = function(x, y, w, h) {
		        grid.rect(x, y + h*0.45 , w-unit(0.5, "mm"), 0.44* h, gp = gpar(fill = col["Clonal"], col = NA), just = "top")
		    },
		    Subclonal = function(x, y, w, h) {
		        grid.rect(x, y + h*0.45 , w-unit(0.5, "mm"), 0.44* h, gp = gpar(fill = col["Subclonal"], col = NA), just = "top")
		    },
		    LOH = function(x, y, w, h) {
				grid.rect(x, y - h*0.45 , w-unit(0.5, "mm"), 0.44* h, gp = gpar(fill = col_cnv["LOH"], col = NA), just = "bottom")
		    },
		    AMP = function(x, y, w, h) {
				grid.rect(x, y - h*0.45 , w-unit(0.5, "mm"), 0.44* h, gp = gpar(fill = col_cnv["AMP"], col = NA), just = "bottom")
		    },
		    LOSS = function(x, y, w, h) {
				grid.rect(x, y - h*0.45 , w-unit(0.5, "mm"), 0.44* h, gp = gpar(fill = col_cnv["LOSS"], col = NA), just = "bottom")
		    },
		    show_legend = FALSE 
		)
	}

	alter_fun = list(
	    background = function(x, y, w, h) {
	        grid.rect(x, y, w-unit(0.5, "mm"), h-unit(0.5, "mm"), 
	            gp = gpar(fill = "#F2F2F2", col = NA))
	    },
	    Clonal = function(x, y, w, h) {
	        grid.rect(x, y , w-unit(0.5, "mm"), h-unit(0.5, "mm"), gp = gpar(fill = col["Clonal"], col = NA))
	    },
	    Subclonal = function(x, y, w, h) {
	        grid.rect(x, y , w-unit(0.5, "mm"), h-unit(0.5, "mm"), gp = gpar(fill = col["Subclonal"], col = NA))
	    },
	    show_legend = FALSE 
	)

	## 顶部注释 class
	top_annotation <- HeatmapAnnotation(
		foo = anno_empty(height = unit(2, "cm") , border =F) ,  ## 留给注释空间
		Class = class_order ,
		col = list( 
			Class = col_class 
		) ,
		annotation_name_side = "left" , 
	  	border = T ,
	  	gap = unit(1, "mm") ,
	  	show_annotation_name = c(Mut_Num = FALSE) , 
	  	annotation_name_gp = gpar(fontsize = 12),
	  	show_legend = FALSE  ## 所有的legend均后期定义
	)

	################################################################################################

	ldg_Class = Legend(labels = names(col_class) , title = "Pathogenic" ,
	 legend_gp = gpar(fill = col_class , bar_width = 1 , fontsize = 12) , ncol = 3 , 
	 gap = unit(1, "cm") ## 图例的间隔
	)

	col_Mut <- col[1:2]
	ldg_Variant = Legend(labels = names(col_Mut) , title = "Clonal status" , border = "black" , ## 注释边框加黑
	 	legend_gp = gpar(fill = col_Mut , bar_width = 1 ,fontsize = 12) , ncol = 3 , 
	 	gap = unit(1, "cm") ## 图例的间隔
	)

	col_Mut <- col 
	ldg_CNV = Legend(labels = names(col_cnv) , title = "CNV status" , border = "black" , ## 注释边框加黑
	 	legend_gp = gpar(fill = col_cnv , bar_width = 1 ,fontsize = 12) , ncol = 3 , 
	 	gap = unit(1, "cm") ## 图例的间隔
	)	

	#lgd_all <- packLegend(ldg_Class , ldg_Variant , ldg_CNV , column_gap = unit(1, "cm") , row_gap = unit(10, "mm") , direction = "horizontal"  )
	lgd_all <- packLegend(ldg_Class , ldg_Variant , column_gap = unit(1, "cm") , row_gap = unit(10, "mm") , direction = "horizontal"  )

	################################################################################################
	p <- oncoPrint(mut, name = "cases", ## 后面加分割线用的name
	    alter_fun = alter_fun, col = col, 
	    top_annotation = top_annotation  ,
	    row_names_side = "left", row_names_gp = gpar(fontsize = 16) ,  ## 基因名移到左边
	    left_annotation = NULL , right_annotation = NULL ,  ## 基因在样本间的柱状分布图
	    show_pct = FALSE , ##不展示百分比
	    border = TRUE,
	    row_order = 1:nrow(mut) ,
	    column_order = sample_order,
	    column_split = type_order , column_title = NULL , ## 按照class拆分成几块
	    show_heatmap_legend = FALSE
	)

	width = 10
	height = 2.2

	##
	pdf(images_name , width = width , height = height)
	draw(p )
	#draw(ldg_Variant, x = unit(0.3, "npc"), y = unit(0.005, "npc"), just = c("left", "bottom"))
	draw(lgd_all, x = unit(0.3, "npc"), y = unit(0.99, "npc"), just = c("left", "top"))
	# draw(ldg_Variant, x = unit(0.99, "npc"), y = unit(0.99, "npc"), just = c("right", "top"))

	################################################################################################
	#### 癌前和癌加分割线
	## 主体
	## 每个人的所有样本用线分割
	if(1!=1){
		for(sample in info[info$Class_sub=="IM-1","Tumor"][-1]){

			decorate_heatmap_body("cases", {
			    i =  which(colnames(mut) %in% sample )[1] -1 
			    x = i/ncol(mut)
			    grid.lines(c(x, x), c(0, 1), gp = gpar(lwd = 1.3, lty = 3)) ## lwd代表粗细，lty代表线的类型
			})

			## 注释在对应位置加分割线
			for(anno in annotation_name){
				decorate_annotation(c(anno ), {
				    i =  which(colnames(mut) %in% sample )[1] -1 
				    x = i/ncol(mut)
				    grid.lines(c(x, x), c(0, 1), gp = gpar(lwd = 1.3 , lty = 3)) ## lwd代表粗细，lty代表线的类型
				})
			}
		}
	}

	dev.off()
}


###########################################################################################
## 合并所有的MAF文件
dat <- dat_ccf
dat$Tumor <- dat$Sample

## 产生输入矩阵
Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
Variant_Type_Combine <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift","In_Frame","Splice_Site","Nonstop_Mutation","Multiple_Hits")

MutMatrix <- CreateMutMatrix( dat = dat , Variant_Type = Variant_Type )

## 瀑布图
images_name <- paste0(images_path , "/Mut_WaterFall.SortyBySample.Driver.pdf")
plotMutWaterFull( MutMatrix = MutMatrix , col_class = col_class , Variant_Type_Combine = Variant_Type_Combine , images_name = images_name )


